package tech.waterism.topo.calculator;

import lombok.extern.slf4j.Slf4j;
import tech.waterism.topo.context.CalculationContext;
import tech.waterism.topo.model.*;

import java.util.*;

/**
 * 预报计算器
 * 
 * <p>负责汇流计算，将上游来水汇总到当前断面
 * 
 * <h3>核心职责：</h3>
 * <ul>
 *   <li>整理上游输入分量</li>
 *   <li>累加入流（QSim）</li>
 *   <li>叠加基流</li>
 *   <li>处理枢纽分流</li>
 * </ul>
 * 
 * @author waterism
 * @since 1.0.0
 */
@Slf4j
public class ForecastCalculator {
    
    /** 河道直接入流类型集合 */
    private static final Set<String> DIRECT_INFLOW_TYPES = Set.of("D", "R");
    
    /** 区间面积入流类型 */
    private static final String AREA_INFLOW_TYPE = "X";
    
    /** 水库入流类型 */
    private static final String RESERVOIR_INFLOW_TYPE = "RSVR";
    
    /** 蓄滞洪区入流类型 */
    private static final String STORE_INFLOW_TYPE = "STORE";
    
    /** 分流闸门类型 */
    private static final String DISTRIBUTION_GATE_TYPE = "DIS";
    
    /**
     * 计算节点的预报结果
     * 
     * @param node 拓扑节点
     * @param context 计算上下文
     * @return 预报结果（包含累加后的QSim）
     */
    public NodeResult calculate(TopoNode node, CalculationContext context) {
        List<InputComponent> inputs = node.getInputs();
        
        // 如果没有输入分量，返回空结果
        if (inputs == null || inputs.isEmpty()) {
            log.debug("节点 {} 没有输入分量，跳过预报计算", node.getStcd());
            return new NodeResult(node.getStcd(), node.getData());
        }
        
        log.debug("开始计算节点 {} 的预报结果，输入分量数：{}", node.getStcd(), inputs.size());
        
        // 1. 整理输入分量
        prepareInputComponents(inputs, context);
        
        // 2. 累加入流
        sumInflows(node, inputs);
        
        // 3. 叠加基流
        addBaseFlow(node);
        
        // 4. 处理枢纽分流
        if (node.isHubNode()) {
            processHubDistribution(node, context);
        }
        
        log.debug("节点 {} 预报计算完成", node.getStcd());
        return new NodeResult(node.getStcd(), node.getData());
    }
    
    /**
     * 整理输入分量
     * <p>从上游节点获取已计算的结果，并根据入流类型计算当前节点的入流
     * 
     * @param inputs 输入分量列表
     * @param context 计算上下文
     */
    private void prepareInputComponents(List<InputComponent> inputs, CalculationContext context) {
        for (InputComponent input : inputs) {
            // 跳过没有模型配置的分量
            if (input.getMdcdList() == null || input.getMdcdList().isEmpty()) {
                log.debug("输入分量 {} 没有模型配置，跳过", input.getArcd());
                continue;
            }
            
            // 从缓存获取上游节点的计算结果
            NodeResult upstreamResult = context.getCachedResult(input.getArcd());
            if (upstreamResult == null) {
                log.warn("未找到上游节点 {} 的计算结果，跳过", input.getArcd());
                continue;
            }
            
            // 根据入流类型计算流量
            List<TimeSeriesData> upstreamData = upstreamResult.getData();
            calculateInflowByType(input, upstreamData, context);
            
            // 处理枢纽分流：如果上游是枢纽，使用枢纽分流数据
            if (context.hasHubDistribution(input.getArcd())) {
                applyHubDistribution(input, context);
            }
        }
    }
    
    /**
     * 根据入流类型计算流量
     * 
     * @param input 输入分量
     * @param upstreamData 上游时序数据
     * @param context 计算上下文
     */
    private void calculateInflowByType(InputComponent input, List<TimeSeriesData> upstreamData, 
                                      CalculationContext context) {
        List<TimeSeriesData> inputData = input.getData();
        String artp = input.getArtp();
        
        for (int i = 0; i < inputData.size() && i < upstreamData.size(); i++) {
            TimeSeriesData current = inputData.get(i);
            TimeSeriesData upstream = upstreamData.get(i);
            
            // 只计算预报时段（在起调时间之后）
            if (!current.isAfter(context.getStartTime())) {
                // 实测时段，使用实测流量
                if (current.getQ() == null) {
                    current.setQ(0.0);
                }
                continue;
            }
            
            // 根据入流类型计算
            Double flow = 0.0;
            
            if (DIRECT_INFLOW_TYPES.contains(artp)) {
                // 河道直接入流：考虑分洪排洪影响
                flow = calculateDirectInflow(upstream);
            } else if (AREA_INFLOW_TYPE.equals(artp)) {
                // 区间面积入流
                flow = Optional.ofNullable(upstream.getQSim()).orElse(0.0);
            } else if (RESERVOIR_INFLOW_TYPE.equals(artp)) {
                // 水库出流
                flow = Optional.ofNullable(upstream.getQOut()).orElse(0.0);
            } else if (STORE_INFLOW_TYPE.equals(artp)) {
                // 蓄滞洪区入流
                flow = Optional.ofNullable(upstream.getQSim()).orElse(0.0);
            }
            
            // 扣除取水量
            flow -= calculateIntakes(input.getIntakes(), i);
            
            current.setQ(flow);
        }
    }
    
    /**
     * 计算河道直接入流
     * <p>考虑分洪排洪影响：QSim + QDsc - QDvr
     * 
     * @param upstream 上游时序数据
     * @return 计算后的流量
     */
    private Double calculateDirectInflow(TimeSeriesData upstream) {
        Double qSim = Optional.ofNullable(upstream.getQSim()).orElse(0.0);
        Double qDsc = Optional.ofNullable(upstream.getQDsc()).orElse(0.0);
        Double qDvr = Optional.ofNullable(upstream.getQDvr()).orElse(0.0);
        return qSim + qDsc - qDvr;
    }
    
    /**
     * 计算取水量
     * 
     * @param intakes 取水口列表
     * @param index 时间序列索引
     * @return 总取水量
     */
    private Double calculateIntakes(List<IntakeInfo> intakes, int index) {
        if (intakes == null || intakes.isEmpty()) {
            return 0.0;
        }
        
        double totalIntake = 0.0;
        for (IntakeInfo intake : intakes) {
            if (intake.getData() == null || intake.getData().isEmpty()) {
                continue;
            }
            if (index < intake.getData().size()) {
                Double q = intake.getData().get(index).getQ();
                totalIntake += Optional.ofNullable(q).orElse(0.0);
            }
        }
        return totalIntake;
    }
    
    /**
     * 应用枢纽分流数据
     * 
     * @param input 输入分量
     * @param context 计算上下文
     */
    private void applyHubDistribution(InputComponent input, CalculationContext context) {
        List<TimeSeriesData> distributionData = context.getHubDistribution(input.getArcd());
        if (distributionData == null) {
            return;
        }
        
        // 将分流数据映射为时间->数据的映射
        Map<Date, TimeSeriesData> distributionMap = new HashMap<>();
        for (TimeSeriesData data : distributionData) {
            distributionMap.put(data.getDt(), data);
        }
        
        // 应用分流数据到输入分量
        List<TimeSeriesData> inputData = input.getData();
        for (TimeSeriesData current : inputData) {
            if (!current.isAfter(context.getStartTime())) {
                continue;
            }
            
            TimeSeriesData distribution = distributionMap.get(current.getDt());
            if (distribution != null) {
                // 根据入流类型选择对应的字段
                if (RESERVOIR_INFLOW_TYPE.equals(input.getArtp())) {
                    current.setQ(Optional.ofNullable(distribution.getQOut()).orElse(0.0));
                } else {
                    current.setQ(Optional.ofNullable(distribution.getQSim()).orElse(
                            Optional.ofNullable(distribution.getQ()).orElse(0.0)));
                }
            }
        }
    }
    
    /**
     * 累加所有输入分量的入流，得到断面总入流QSim
     * 
     * @param node 拓扑节点
     * @param inputs 输入分量列表
     */
    private void sumInflows(TopoNode node, List<InputComponent> inputs) {
        List<TimeSeriesData> nodeData = node.getData();
        
        // 初始化：将所有QSim清零
        for (TimeSeriesData data : nodeData) {
            data.setQSim(0.0);
        }
        
        // 累加所有输入分量
        for (InputComponent input : inputs) {
            List<TimeSeriesData> inputData = input.getData();
            for (int i = 0; i < nodeData.size() && i < inputData.size(); i++) {
                Double inputQSim = Optional.ofNullable(inputData.get(i).getQSim()).orElse(0.0);
                nodeData.get(i).addQSim(inputQSim);
            }
        }
        
        log.debug("节点 {} 累加入流完成，时段数：{}", node.getStcd(), nodeData.size());
    }
    
    /**
     * 叠加基流（地下水补给）
     * 
     * @param node 拓扑节点
     */
    private void addBaseFlow(TopoNode node) {
        List<Double> qBases = node.getQBases();
        if (qBases == null || qBases.isEmpty()) {
            return;
        }
        
        List<TimeSeriesData> nodeData = node.getData();
        for (int i = 0; i < qBases.size() && i < nodeData.size(); i++) {
            Double baseFlow = qBases.get(i);
            if (baseFlow != null) {
                nodeData.get(i).addQSim(baseFlow);
            }
        }
        
        log.debug("节点 {} 叠加基流完成", node.getStcd());
    }
    
    /**
     * 处理枢纽分流
     * <p>如果当前节点是枢纽分流口，记录分流数据供下游使用
     * 
     * @param node 拓扑节点
     * @param context 计算上下文
     */
    private void processHubDistribution(TopoNode node, CalculationContext context) {
        HubInfo hub = node.getHub();
        if (hub == null || hub.getGate() == null) {
            return;
        }
        
        // 查找分流闸门
        for (GateInfo gate : hub.getGate()) {
            if (DISTRIBUTION_GATE_TYPE.equals(gate.getGttp())) {
                String distributionStcd = gate.getDtcd();
                context.addHubDistribution(distributionStcd, node.getData());
                log.debug("节点 {} 记录枢纽分流数据到 {}", node.getStcd(), distributionStcd);
                break;
            }
        }
    }
}

